Single gpu disk offload PTQ for DSR1/Ultra - #2008
Conversation
📝 WalkthroughWalkthroughChangesDisk-offloaded Hugging Face model loading now accepts memory caps and offload folders, while unified export adds streaming support for Accelerate-offloaded models. Export state handling is refined for tied caches, meta tensors, state-dict rewriting, and MoE aliases. A new PTQ recipe and validation coverage are included. Offload-aware model loading and export
Export state and tied-weight handling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant get_model
participant Accelerate
participant export_hf_checkpoint
participant StreamingShardWriter
CLI->>get_model: provide offload folder and memory caps
get_model->>Accelerate: load with device map and offload settings
Accelerate-->>get_model: return offloaded model
CLI->>export_hf_checkpoint: export quantized model
export_hf_checkpoint->>StreamingShardWriter: stream processed tensors
StreamingShardWriter-->>export_hf_checkpoint: write shards and index
export_hf_checkpoint-->>CLI: write Hugging Face export configuration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
52fc552 to
827bbe7
Compare
Fridah-nv
changed the title
|
/claude review |
827bbe7 to
8597e97
Compare
8597e97 to
87ccd1a
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2008 +/- ##
==========================================
+ Coverage 60.23% 61.77% +1.54%
==========================================
Files 519 519
Lines 59399 59687 +288
==========================================
+ Hits 35778 36874 +1096
+ Misses 23621 22813 -808
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
RB: Can you explain to me the overall flow of the changes in this file?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 I’m tracing the file-level flow and its before/after behavior now.
There was a problem hiding this comment.
🤖 Bot comment.
The high-level change is that export_hf_checkpoint() now has two Transformers paths:
- Normal model: keeps the existing flow: process all quantized modules → build/post-process one full state dict → call
save_pretrained(). - Accelerate CPU/disk-offloaded model:
_has_accelerate_offload()detects the hooks and routes to_export_transformers_checkpoint_streaming().
The new offloaded flow is:
- Run the same model-level preparation as before (MoE preparation, resmooth/requantize, build quant config, and synchronize relevant amax values).
- Discover decoder layers, then materialize one layer at a time inside
enable_weight_access_and_writeback(..., writeback=False). While the real weights are available, dispatch the existing export handlers to pack quantized weights and reconstruct fused MoE weights. - Immediately post-process each tensor (drop quantizer internals, convert KV-cache amax to scales, squeeze scales, reverse Transformers conversion names, and skip tied aliases), move it to CPU, and feed it to
_StreamingShardWriter. Exiting the context lets Accelerate offload that layer again. - Repeat the same materialize-and-stream treatment for offloaded non-decoder modules such as embeddings, norms, and
lm_head, then stream any remaining resident parameters/buffers. _StreamingShardWriterflushes atmax_shard_size, renames temporary parts to canonical safetensor shard names, and writes the index. This bounds peak CPU memory to roughly one materialized layer plus one shard buffer instead of a full-model state dict.- Call
save_pretrained(state_dict={})only to emit config/custom-code artifacts, then the outer function writeshf_quant_config.jsonand updates/sanitizesconfig.jsonas before.
Two safety pieces support this: Accelerate hooks are no longer removed before offloaded weights are materialized, and _export_quantized_weight() now fails loudly if a meta tensor reaches it. _process_quantized_modules_offloaded() is the analogous full-state-dict helper for direct lower-level calls, but the public top-level offload path uses streaming to avoid the full CPU-RAM accumulation.
There was a problem hiding this comment.
RB: This PR should treat offloading as the special case and consider not offloading as the general case. I feel that this PR does not respect this assumption. Example there are many places where regular path is route to if not _offloaded: ... (is it possible to have a cleaner seperation which does not need to move general model export code under if not _offloaded?)
There was a problem hiding this comment.
🤖 Bot comment.
🐝 I’m checking the branch structure and shared-vs-offload responsibilities now.
There was a problem hiding this comment.
🤖 Bot comment.
Yes—your read is correct. The regular path should remain the straight-line default, with one early special-case dispatch:
if _has_accelerate_offload(model):
return _export_offloaded_transformers_checkpoint(...)
# existing regular export flow, unchanged
post_state_dict, hf_quant_config = _export_transformers_checkpoint(...)
...
model.save_pretrained(...)The offload helper should own its unsupported-option warnings, layer materialization/streaming, inline name conversion, and save_pretrained(state_dict={}) artifact workaround. Only genuinely common finalization (quant config + sanitized config.json) needs a small shared helper.
There is also duplicated dispatch today: export_hf_checkpoint() routes offloaded models directly to _export_transformers_checkpoint_streaming(), while _export_transformers_checkpoint() independently detects _offloaded and branches to _process_quantized_modules_offloaded(). That lower-level branch is bypassed by the public path and makes ordinary callers offload-aware. Removing that branch from the general exporter and keeping it inside the dedicated offload path would restore the intended separation and eliminate the repeated if not _offloaded guards.
There was a problem hiding this comment.
@Fridah-nv could you please address my concern - https://github.com/NVIDIA/Model-Optimizer/pull/2008/changes#r3648000772
Here is bot comment: https://github.com/NVIDIA/Model-Optimizer/pull/2008/changes#r3648016031
There was a problem hiding this comment.
Addressed in f762313: the offloaded path now early-returns after its own block, and a _write_hf_export_config helper owns the shared hf_quant_config.json + config.json finalization. The regular export path is now the straight-line default with no guards; the offloaded block is a clean early exit at the top.
| _sanitize_generation_config_for_save(model) | ||
| _patches = _patch_revert_weight_conversion() | ||
| try: | ||
| model.save_pretrained(str(export_dir), state_dict={}) |
There was a problem hiding this comment.
It seems this line may remove existing safetensors files because save_pretrained with an empty state_dict will think the already written safetensors files obsolete and so remove them. could you double check?
There was a problem hiding this comment.
This is correct. I forgot to push the local fix e4d69c7. It calls with explicit model.config.save_pretrained() + generation_config.save_pretrained() + shutil.copy2(*.py). This also avoids the MoE shared-tensor crash we hit on DSR1, where save_pretrained fired safetensors' shared-tensor validation on the live model parameters even with an empty state dict.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Inline notes on the five items I consider merge-blocking or worth resolving before merge. Overall the streaming-export architecture is sound and the Ultra-550B validation run (161 GB peak VRAM, 15 shards) is convincing evidence the happy path works — these are about CI, an aliasing edge case, and two scope questions.
| self._buffer = {} | ||
| self._buffer_bytes = 0 | ||
|
|
||
| def add(self, key: str, tensor: torch.Tensor) -> None: |
There was a problem hiding this comment.
The streaming path drops the tied-weight safety net, so save_file can crash.
The batch path ends postprocess_state_dict with a data_ptr()-based dedup (quant_utils.py ~L1115-1135) that removes any aliased tensor. That's what guarantees safetensors.save_file never sees two keys sharing storage. The streaming path replaces it with name matching against _tied_weights_keys, gated on tie_word_embeddings=True.
_stream_tensor passes tensor.detach().contiguous().cpu() — for an already-contiguous CPU tensor all three are no-ops returning the same object, so aliasing survives into self._buffer and save_file raises RuntimeError: Some tensors share memory.
Concrete failure: a model with tie_word_embeddings=False in config but genuinely shared lm_head.weight / embed_tokens.weight storage, or any aliasing not literally listed in _tied_weights_keys (note transformers v5 changed _tied_weights_keys to a dict of patterns for several architectures, so exact-string matching can silently miss). Both keys land in the same shard buffer → crash after all prior shards are already written.
The comment above is right that data_ptr() is unreliable across the whole export — but it is reliable within a single shard buffer, because the writer holds a strong reference to every buffered tensor until _flush(). A per-buffer check restores parity without the cross-layer staleness problem:
def add(self, key: str, tensor: torch.Tensor) -> None:
ptr = tensor.data_ptr()
if ptr in self._buffer_ptrs: # reset alongside self._buffer in _flush()
return # (or .clone() if you'd rather keep both keys)
self._buffer_ptrs[ptr] = key
self._buffer[key] = tensor
...There was a problem hiding this comment.
I adopted the suggested change. In the streaming export case, the tied word embedding case is handled by name filter in _stream_tensor, while the storage need to guard on alias because of these cases:
ExportContext.tied_cache / moe_tied_cache are live dicts, and _export_quantized_weight deliberately re-points weight / weight_scale / weight_scale_2 / input_scale at a previously-processed module sharing source memory — so the downstream data_ptr dedup can collapse them. This is postprocess_state_dict, batch-path only; streaming has no equivalent, so such aliases would reach save_file directly.
Let me know if this makes sense. I'm actually not famliar with the moe_tied_cache part.
| f"Architecture {architecture} not found in transformers: {transformers.__version__}. " | ||
| "Falling back to AutoModelForCausalLM (or AutoModel for non-causal architectures)." | ||
| ) | ||
| if not hasattr(transformers, architecture): |
There was a problem hiding this comment.
Unexplained behavior change: dropping or "Deepseek" in architecture from this condition.
Previously any DeepSeek architecture took the AutoModelForCausalLM + assert trust_remote_code path even when transformers exposed the class. Now DeepseekV3ForCausalLM (present since transformers 4.53) falls through to getattr(transformers, architecture)._from_config, and further down model_kwargs2.pop("trust_remote_code", None) fires.
That's a real behavior change for the exact model family this PR targets, it isn't mentioned in the description, and there's no test covering it. It also appears to pull in the opposite direction from _install_transformers_compat_shims(), which exists specifically to make remote-code DSR1 load.
Could you explain the intent, or split this into its own PR so it can be evaluated independently?
There was a problem hiding this comment.
I clean this up as well but need a full run on DSR1 to verify
02d9ac0 to
1ada143
Compare
6aad032 to
6538eb1
Compare
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Re-reviewed at e519853 — all five of my earlier findings are resolved, and I verified the first two by running the checks (ruff format/ruff check clean; recipe docs now 21/21 with no stale rows). One remaining test-coverage gap below.
| ) | ||
|
|
||
| # 2. All tensors in safetensors shards must be non-empty (no meta serialized as zeros) | ||
| safetensor_files = list(export_dir.glob("*.safetensors")) |
There was a problem hiding this comment.
Missing: a streaming-vs-batch key-parity test.
The unit tests added for the alias guard and the data_ptr-under-offload cases are good, and this test is a solid guard against the original meta-tensor bug. But every assertion here iterates only the keys that were written — so it can confirm the written tensors are sane, and can't detect anything missing.
Concretely: if a refactor caused _export_transformers_checkpoint_streaming to skip an entire decoder layer (a seen_keys collision, a _tied_weights_keys pattern over-matching and dropping too much, a _dispatch_export_handler early-return, or the non-decoder pass missing an offloaded lm_head), this test still passes — the surviving shards are non-empty and non-meta.
That's the highest-risk failure mode for this PR, because the streaming path reimplements a lot of _export_transformers_checkpoint + postprocess_state_dict per-tensor rather than sharing it, and the two can drift silently.
The cheap version is to export the same tiny model twice — once offloaded, once not — and diff the key sets:
def test_streaming_export_matches_batch_export_keys(tmp_path):
# non-offloaded reference
ref_model = ... # same tiny llama, plain .cuda(), same quant_cfg + forward_loop
export_hf_checkpoint(ref_model, export_dir=str(ref_dir))
# offloaded, exercises the streaming path
off_model, *_ = _make_cpu_offloaded_model(...)
export_hf_checkpoint(off_model, export_dir=str(off_dir))
assert _all_keys(off_dir) == _all_keys(ref_dir)where _all_keys unions safe_open(...).keys() across every shard. Comparing values too would be even better, but key parity alone catches the whole "silently dropped tensors" class and is what the current assertions can't see.
Two smaller gaps in the same spirit, if you want them in the same pass:
- Multi-shard is never integration-tested. Tiny LLaMA fits in one shard, so
model.safetensors.index.jsongeneration inside the real export path is untested end-to-end — only_StreamingShardWriteris tested in isolation. Parametrizing this test withmax_shard_size="1MB"would cover it. _parse_shard_sizehas no test.
There was a problem hiding this comment.
You need to also enable tools unit tests in https://github.com/NVIDIA/Model-Optimizer/blob/main/.github/workflows/unit_tests.yml
31f8164 to
0da6b56
Compare
### What does this PR do?
**Type of change:** New feature (developer tool) + new tests
Adds a **standalone, cross-process resource monitor** —
`tools/resource_monitor.py` —
that samples a workload's GPU and CPU usage *from the outside* while it
runs, so
we can verify per-run memory budgets (e.g. the single-GPU layerwise PTQ
target of
≤80 GB GPU / ≤80 GB CPU for OMNIML-4947) without instrumenting
`hf_ptq.py` itself.
The monitor wraps any command, samples at a fixed interval, and on exit
writes a
CSV timeseries plus a peak/mean/min summary:
- **GPU** (per device, via NVML with an `nvidia-smi` fallback): used
memory,
utilization %, power draw (W), temperature (°C).
- **CPU** (via `psutil`): system total/used/free memory + utilization %,
and the
monitored process tree's RSS + CPU %.
It is opt-in from `examples/hf_ptq/scripts/huggingface_example.sh` via
`MODELOPT_MEM_MONITOR=1` (off by default → byte-for-byte identical
behavior). When
enabled it wraps the `hf_ptq.py` run and writes the trace/summary to a
**sibling**
`${SAVE_PATH}_mem_monitor/` directory, kept out of the exported
checkpoint that is
uploaded and consumed downstream.
**Files:**
- `tools/resource_monitor.py` — the sidecar (NVML + `nvidia-smi`
fallback; `psutil`).
- `tests/unit/tools/test_resource_monitor.py` — CPU-only unit tests (run
in the `unit` nox lane).
- `examples/hf_ptq/scripts/huggingface_example.sh` — opt-in
`MODELOPT_MEM_MONITOR=1` wrapper.
- `examples/hf_ptq/requirements.txt` — adds `psutil`.
- `pyproject.toml` — adds `psutil` to the `dev-test` extra
(deterministic import in the unit lane).
- `.github/workflows/unit_tests.yml` — adds `tools/resource_monitor.py`
to the unit-test path filters.
#### Why a new tool instead of extending
`modelopt/torch/utils/memory_monitor.py`?
The existing `GPUMemoryMonitor` is a fundamentally different tool and
cannot serve
this use case by extension:
| | `modelopt.torch.utils.memory_monitor.GPUMemoryMonitor` |
`tools/resource_monitor.py` (this PR) |
|---|---|---|
| Scope | **In-process** thread inside the workload | **Cross-process**
— wraps an external command |
| Survives workload OOM/SIGKILL | ❌ dies with the process | ✅ keeps
sampling, still writes the summary |
| Import cost | Pulls `torch` (~19 s) — lives in the workload |
Torch-free (`psutil`+`pynvml`, ~0.03 s) |
| Metrics | GPU device memory only | GPU mem/util/**power/temp** +
**CPU** mem/util + process-tree RSS |
| Output | In-memory / logs | CSV timeseries + peak/mean/min summary |
Merging the two would force `torch` into a standalone sidecar (defeating
the point)
or split the in-process monitor's threading model. A future refactor may
factor out
a **shared torch-free sampling core with two thin frontends**
(in-process + sidecar);
that is tracked as a follow-up rather than blocking this monitoring
harness, which
PR #2008 (single-GPU disk-offload PTQ) depends on.
### Usage
```bash
# Wrap mode (preferred): monitor exits with the workload's return code
python tools/resource_monitor.py --gpus 2,3 --out mem.csv --summary peak.txt -- \
python hf_ptq.py --pyt_ckpt_path=<model> --qformat=nvfp4 ...
# Opt-in from the HF PTQ example (off by default):
MODELOPT_MEM_MONITOR=1 CUDA_VISIBLE_DEVICES=2,3 CUDA_DEVICE_ORDER=PCI_BUS_ID \
bash examples/hf_ptq/scripts/huggingface_example.sh <args>
# -> writes ${SAVE_PATH}_mem_monitor/mem_trace.csv and mem_peak.txt
```
### Testing
- **Unit (CPU-only, in the `unit` nox lane):** `pytest
tests/unit/tools/test_resource_monitor.py`
— 11 tests covering `--gpus` parsing (CSV + space-separated, UUID/MIG
rejection),
the disabled/`nvidia-smi` sampling paths (including `[N/A]` → `None` and
the
smi-failure-yields-empty guard), CPU sampling, the accumulator, and
end-to-end
CSV/summary + exit-code propagation in wrap mode.
- **GPU-validated** on a B200 node (GPUs 2,3): confirmed the `gpu{i}_*`
memory /
utilization / power / temperature columns populate and the summary is
written.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ — new tool; the example wrapper
is off unless `MODELOPT_MEM_MONITOR=1`.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — `psutil`
added to `dev-test` + the hf_ptq example requirements;
`pynvml`/`nvidia-ml-py` is an optional runtime dep (graceful
`nvidia-smi` fallback).
- Did you write any new necessary tests?: ✅ —
`tests/unit/tools/test_resource_monitor.py`.
- Did you update Changelog?: N/A — repo-level `tools/` script, not
shipped in the wheel.
- Did you get Claude approval on this PR?: ❌ <!-- run /claude review -->
### Additional Information
Part of **OMNIML-4947** (single-GPU disk-offload PTQ). This is PR 1 of
the stack —
the monitoring harness that PR #2008 (disk-offload layerwise PTQ +
offload-aware
export) builds on.
---------
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
3242b08 to
b2d731d
Compare
G1 — offload-aware unified HF export - Add _has_accelerate_offload() to detect CPU/disk-offload hooks - Add _process_quantized_modules_offloaded() that materializes decoder layers one at a time inside enable_weight_access_and_writeback, then collects the full state dict via a second loop for non-decoder offloaded modules (embed, norm, lm_head) - Branch _export_transformers_checkpoint on the offload flag; remove hooks only after the offloaded state dict has been collected - Add meta-tensor guard in _export_quantized_weight to catch accidental standard-path use on offloaded models G2 — disk-offload CLI wiring in hf_ptq - Add --offload_folder, --max_gpu_memory_gb, --max_cpu_memory_gb args - Inject max_memory budget into load_model_from_config; skip seq_device_map when offload folder is set - Add nvfp4_experts_only-kv_fp8_layerwise_offload.yaml recipe Other - Fix _FP8BF16Fallback shim: nested try avoids ambiguous except; remove how-comments from matmul; tighten outer except to Exception only - Fix get_nemotron_h_decoder_layers to check both backbone.layers and model.layers Tests - 7 CPU-only unit tests (test_offload_export.py) - 2 GPU integration tests (tests/gpu/torch/export/test_offload_export.py) Signed-off-by: Fridah-nv <fridah@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
embed_tokens, final norms, and lm_head are disk-offloaded alongside decoder layers on single-GPU runs. model.state_dict() returns meta placeholders for them; after revert_weight_conversion_quant_aware renames to hub-original keys, transformers' save_pretrained looks up the tensors by hub name and crashes (e.g. NemotronHForCausalLM has no attribute 'backbone'). Fix: add a second loop in _process_quantized_modules_offloaded that iterates non-decoder modules, skips any without a live offload hook, checks for DIRECT meta parameters/buffers (avoids re-collecting decoder children already captured above), and materializes each via enable_weight_access_and_writeback. Signed-off-by: Fridah-nv <fridah@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ndler, memory Four correctness / efficiency bugs in _process_quantized_modules_offloaded: 1. MoE reconstruction ordering: _reconstruct_fused_moe_linear was called after _process_quantized_modules_offloaded returned the state dict, so fused-MoE decoder layers shipped with per-expert 2D keys (or meta weights after context exit). Move it per-layer inside the materialization window before the state dict snapshot, mirroring the non-offloaded path. 2. Quantized non-decoder modules missed: lm_head (or any quantized module outside the decoder stack) never had _dispatch_export_handler called, so it exported raw unquantized weights. Add handler dispatch inside the non-decoder materialization context. 3. Decoder snapshots accumulating on GPU: tensor.detach() kept every layer's materialized weights on the GPU, defeating the layer-by-layer memory goal. Change to tensor.detach().cpu(). 4. writeback=True on decoder context: on context exit the quantized weights were written back to the offload store (disk → CPU promotion per layer), wasting CPU memory. Changed to writeback=False since weights are captured in layer_tensors immediately after. Also: apply gpu_mem_percentage (default 80 %) when --max_gpu_memory_gb is not specified in disk-offload mode, matching the documented CLI default. Restore test_non_decoder_offloaded_tensors_are_collected (lost during squash). Signed-off-by: Fridah-nv <fridah@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Replace the accumulate-then-save pattern in the offloaded export path with a stream-then-save pattern. Peak memory drops from ~764 GiB (Ultra 550B full state dict in RAM) to 1 decoder layer + 1 shard buffer (~57 GB). New pieces: - `_postprocess_single_tensor` in quant_utils.py: per-tensor subset of postprocess_state_dict for use in the streaming loop - `_StreamingShardWriter`: buffers tensors up to max_shard_size, flushes to temp part files, renames to canonical shard names at finalize() - `_parse_shard_size`: converts "10GB"/"500MB" strings to bytes - `_export_transformers_checkpoint_streaming`: streams decoder layers one at a time via enable_weight_access_and_writeback, applies per-tensor postprocessing and name reversal inline, handles tied-weight dedup from _tied_weights_keys - `export_hf_checkpoint` dispatch: branches on _has_accelerate_offload to call the streaming path instead of _export_transformers_checkpoint for offloaded models; hf_quant_config.json and config.json update are shared between paths Unit tests: 10 new tests covering _StreamingShardWriter (single-shard, multi-shard, readback) and _postprocess_single_tensor (passthrough, filter, rename, squeeze, real-quant drop, kv scale divide). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…eview - Extract _KV_CACHE_REPLACEMENTS / _QLORA_REPLACEMENTS / _BASE_SKIP_KEYS / _QLORA_SKIP_KEYS as module-level constants; eliminate the per-call rebuild in both _postprocess_single_tensor and postprocess_state_dict. - Add _maybe_squeeze_scale helper; remove three identical inline squeeze expressions. - Replace _StreamingShardWriter._part_bytes list (used only for sum()) with a scalar _total_bytes accumulator. - Collapse the two GPU-resident named_parameters / named_buffers loops into a single itertools.chain loop. - Move import contextlib to module level (was deferred inside function body). - Initialize export_state_dict = None before the _offloaded branch to prevent a latent NameError on future edits. - Narrow except (ImportError, Exception) to except ImportError in _parse_shard_size so parser errors propagate instead of silently falling through to the manual fallback. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
- Gate tied-weight dedup on model.config.tie_word_embeddings to avoid
incorrectly dropping lm_head.weight when embeddings are not tied
- Add _is_persistent_buffer helper; filter named_buffers() in the
GPU-resident pass to match state_dict() semantics
- Add .contiguous() before .cpu() in _stream_tensor to handle
non-contiguous views from accelerate writeback
- Copy trust_remote_code modeling files via model.save_pretrained(
state_dict={}) with single-shard rename protection so custom model
class files are not lost
- Refactor example_utils.py shims: module-level _FP8BF16Fallback class,
_install_transformers_compat_shims() called lazily from get_model(),
explicit UserWarning for lossy BF16 fallback path
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
model.save_pretrained(state_dict={}) triggers safetensors' shared-tensor
validation on the live model parameters even when no weights are being
saved. DSR1 and other MoE models have expert weights that share storage
across layers, so this check always fails — crashing the export after
all shards are correctly written and leaving hf_quant_config.json unwritten.
Replace with targeted saves:
- model.config.save_pretrained() for config.json
- model.generation_config.save_pretrained() for generation_config.json
- shutil.copy2(*.py) for trust_remote_code custom modeling files
Also add missing contextlib and shutil stdlib imports.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…, shard regression test
- Delete _process_quantized_modules_offloaded: dead from export_hf_checkpoint
(public path dispatches to streaming writer; function accumulated full state dict
in RAM, defeating offload)
- Replace if _offloaded: branch inside _export_transformers_checkpoint with
NotImplementedError — streaming path owns that case via export_hf_checkpoint
- Add _write_hf_export_config helper (hf_quant_config.json + config.json patching)
- Refactor export_hf_checkpoint: early return after offloaded path eliminates
three scattered if not _offloaded: guards; both paths share the helper
- Update meta-guard error message to point to export_hf_checkpoint
- Tests: remove test_non_decoder_offloaded_tensors_are_collected (tests deleted fn);
add test_multi_shard_files_exist_after_finalize (regression: save_pretrained(state_dict={})
triggered transformers cleanup loop that deleted model-NNNNN-of-NNNNN shards)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Addresses PR review: recipe-doc parity, ruff, and mypy all fail on this branch. - modelopt_recipes/ptq.md: add the nvfp4_experts_only-kv_fp8_layerwise_offload row and bump the summary count to 21, restoring the parity enforced by tests/unit/recipe/test_recipe_docs.py. - Apply ruff format to example_utils.py and quant_utils.py. - test_offload_export.py: wrap st.keys() in list() to clear SIM118. Ruff's own suggested fix (iterating the handle directly) would break the test -- safetensors safe_open handles are not iterable. - Fix 5 pre-existing mypy errors: duplicate raw_tied_keys annotation, missing None guard on the _postprocess_single_tensor result (its documented contract is to return (None, None) to skip), rebinding hf_quant_config to a different type in _write_hf_export_config, and two stale type: ignore comments. No behavior change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…s opt-in Addresses review feedback on the streaming export path and the hf_ptq shims. _StreamingShardWriter.add now resolves shared storage before buffering. safetensors.save_file rejects a dict holding two tensors that share memory, and _stream_tensor's .detach().contiguous().cpu() is a no-op for an already-contiguous CPU tensor, so the name-based _tied_weights_keys filter was the only guard -- it misses ties transformers does not declare. data_ptr() is unreliable across the export but reliable within one buffer, since buffered tensors stay alive until flush. Exact (pointer, shape, dtype) matches are dropped as genuine ties; partial matches are cloned so a distinct view is never silently lost. _install_transformers_compat_shims is now opt-in via --allow_compat_shims. It was running on every get_model() call, applying process-wide monkeypatches -- silently disabling flash attention and substituting a lossy BF16 FP8 matmul -- for users who never asked. Gating on trust_remote_code would be wrong: the FP8 shim patches the native HF path, which is what DeepSeek-R1 now uses. The suppress(Exception) around the FP8 block is narrowed so a missing _load_finegrained_fp8_kernel warns instead of silently no-opping; on transformers 5.7 that symbol is absent, so the fallback was never actually installed. _FP8BF16Fallback.matmul now expands scales by block_size rather than out_f // nb_out. The scale grid is ceil-divided, so the ratio mis-groups the last block when a dimension is not a multiple of block_size (out_f=300, block=128 gave 100) and truncates outright when the ratio does not divide evenly. Verified against a per-block reference for divisible and non-divisible shapes. Also drops the fp32 intermediate (~528 MB per call at DSR1's 7168x18432) by scaling in bf16. Per review, hoist the offloaded refusal in _export_transformers_checkpoint to a guard clause beside the detection, removing the else-after-raise and the if not _offloaded wrapper around hook removal. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ote_code The removal of the "Deepseek" clause from the architecture dispatch was an undocumented drive-by in dd76c1b. It had two side effects beyond its intent: --trust_remote_code was silently ignored for the model class, and DeepseekV2 and DeepseekVL were swept onto the built-in path alongside DeepseekV3 despite neither being validated here. Restore the clause, gated on the flag. --trust_remote_code now selects the bundled modeling code as before; without it the built-in class is used, which is what the disk-offload and streaming-export paths are validated against and which previously raised AssertionError. Non-DeepSeek architectures are untouched. Also reword the compat-shim docstring: two of its three patches target bundled modeling files, not the built-in path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ironment These monkeypatched transformers process-wide to paper over environment problems, which does not belong in an example -- especially a quantization one. - _FP8BF16Fallback substituted a lossy BF16 dequant for the block-scaled FP8 matmul, degrading the very forward pass used for calibration amax collection. It also patched _load_finegrained_fp8_kernel, a private symbol that no longer exists in transformers 5.7 (now lazy_load_kernel), so within the supported range (>=4.56,<5.13) it silently no-ops. The fix is `pip install kernels`. - The flash_attn shim forced availability checks to False when the package was installed but unimportable. Uninstalling or repairing the broken install reaches the same state without a global patch. - The is_torch_fx_available shim covered old bundled modeling files on transformers 5; not worth a process-wide patch on its own. Also removes --allow_compat_shims, added in the previous commit to gate them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The offloaded export path used data_ptr() as a stable tensor identity, but it only identifies a tensor while that tensor is resident. Two failure modes, both silent, both verified on Qwen3.6-35B-A3B with a 30/20 GB budget: Recycled addresses. ExportContext.tied_cache / moe_tied_cache persisted for the whole export while the streaming path materialises and frees one module at a time, so the allocator handed a later expert the address of a freed earlier one. The alias step then re-pointed its weight and scales at the wrong module. 1536 false hits, leaving 60% of expert weights (18440 tensors) as byte-identical copies of unrelated experts, and short-circuiting half the export (_export_quantized_weight ran 14592 times instead of 30720). Null addresses. sync_tied_input_amax groups by weight data_ptr, and every meta tensor reports 0, so all offloaded modules collapsed into two buckets whose amaxes were max-merged model-wide. 16896 expert input_scale values inflated, median 90% relative error, worst 12x. Fixes: ExportContext.reset_tied_caches() drops dedup state at the end of each materialization window, and sync_tied_input_amax skips modules whose weights are not resident, warning only when the model declares ties it could not honour. Verified against the batch export path, which is unaffected (it keeps every weight resident): all 123846 shared tensors now match bitwise, versus 70580 differing before. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The streaming path warned that extra_state_dict "is not supported" and dropped it. That silently lost the MTP weights: HF builds only num_hidden_layers decoders, so multi-token-prediction tensors are orphaned and reach export solely through extra_state_dict (hf_ptq.py passes load_mtp_weights' output there). Exporting Qwen3.6-35B-A3B produced 19 fewer tensors than the batch path, all mtp.* -- mtp.fc.weight, mtp.norm.weight, and the whole mtp.layers.0 block including its fused experts. A checkpoint missing them cannot serve speculative decoding. Feed them to the shard writer after the resident-tensor pass. They are already materialized and skip per-tensor postprocessing, matching the batch path which merges them after postprocess_state_dict; only the hub-name reversal applies. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ross-layer aliasing Per-expert weight wrappers inside _export_fused_experts are freshly allocated .contiguous() copies. Their data_ptr() is ephemeral: it is freed when packing replaces the tensor, so the GPU allocator can recycle that address for an unrelated expert in a later layer. Passing those ephemeral addresses into the caller-owned _tied_cache caused false-positive cache hits, silently aliasing the packed weight/scale tensors of one MoE layer onto another (manifesting as 35 k weight_scale mismatches between layers 15-60 in the DSR1 NVFP4 export). Changes: - Drop _tied_cache param from _export_fused_experts signature and all call sites (_export_fused_experts_module in hf_export_handlers.py). - Extend docstring to explain WHY _tied_cache is intentionally excluded. - Update test_fused_experts.py: remove _tied_cache from both dedup test call sites (they only need _moe_tied_cache for module-level dedup, which is unaffected). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
main disables pointer-keyed dedup under FSDP2 by setting ExportContext.tied_cache and moe_tied_cache to None, for the same reason this branch scopes them per materialization window: FSDP2 recycles data_ptr() values as modules are resharded. reset_tied_caches() called .clear() unconditionally and would raise AttributeError on an FSDP2 export, so make it a no-op when dedup is already disabled. Also fixes two artifacts of the rebase resolution: a missing blank line after the FSDP2 helpers in example_utils.py, and an ambiguous multiplication sign in a comment that RUF003 rejects. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 7
🧹 Nitpick comments (3)
modelopt/torch/export/unified_export_hf.py (2)
1277-1283: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRaise the offload rejection before mutating the model.
PrepareMoEInputsRegistrydispatch andrequantize_resmooth_fused_llm_layers(model)(Line 1261-1275) already ran by the time thisNotImplementedErrorfires, so a caller such as a speculative exporter is left with a partially-prepared model and cannot retry cleanly. Move the check to the top of the function, right after the dtype handling.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 1277 - 1283, Move the _has_accelerate_offload(model) validation to the top of _export_transformers_checkpoint, immediately after dtype handling and before PrepareMoEInputsRegistry dispatch or requantize_resmooth_fused_llm_layers(model). Preserve the existing NotImplementedError message and streaming-export guidance, and remove the later duplicate check.
877-889: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTemp shard parts leak into
export_dirif the export aborts.
__shard_part_*.safetensorsfiles are only renamed infinalize(). Any failure between the first flush and finalize (e.g. the meta guard at Line 575, an unsupported MoE handler) leaves them in the user-visible export directory, where they also match*.safetensorsglobs. Consider aclose()/context-manager that unlinks unfinalized parts, invoked from atry/finallyaround the streaming loops.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 877 - 889, The shard writer must clean up unfinalized temporary parts when export fails. Add a close or context-manager cleanup path for the writer containing _flush that unlinks every path in _part_files unless finalize completed, and invoke it from a try/finally surrounding the streaming export loops so failures before finalize cannot leave __shard_part_*.safetensors files in the export directory.modelopt/torch/export/quant_utils.py (1)
962-1049: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider having
postprocess_state_dictdelegate to_postprocess_single_tensor.The new helper documents itself as "Per-tensor subset of :func:
postprocess_state_dict, for streaming export." butpostprocess_state_dict's main loop (L1070-1102) still re-implements the same skip/replace/KV-cache logic independently, and the RealQuant-scale and LoRA-key removal are done as separate post-passes there (L1106-1119) instead of the early-return checks already present in_postprocess_single_tensor(L1017-1023). The two now share the lookup tables, but the control flow itself is duplicated and can silently diverge if either is updated without the other.♻️ Sketch: delegate per-key logic to `_postprocess_single_tensor`
post_state_dict = {} for key, value in state_dict.items(): - # Skip problematic parameters for specific model architectures, e.g., Nemotron Nano VL models - if key == "vision_model.radio_model.summary_idxs": - logger.info(f"Removing problematic parameter: {key}") - continue - - # Skip keys not related to quantizers - if all(skip_key not in key for skip_key in skip_keys): - post_state_dict[key] = value - continue - - # Apply replacements if the key matches any suffix in the replacements dict - for old_suffix, new_suffix in replacements.items(): - if key.endswith(old_suffix): - prefix = key[: -len(old_suffix)] - if "_amax" in key: - ... - post_state_dict[prefix + new_suffix] = value - break - - post_state_dict = {k: _maybe_squeeze_scale(k, v) for k, v in post_state_dict.items()} - - # remove real quant parameters from the state dict - keys_to_delete = [] - for key, value in post_state_dict.items(): - if any( - key.endswith("weight_quantizer." + q_key) - for q_key in RealQuantLinear.list_of_scale_tensors - ): - keys_to_delete.append(key) - - # remove LoRA adapters from state dict - if is_modelopt_qlora: - for key in post_state_dict: - if "lora" in key and key not in keys_to_delete: - keys_to_delete.append(key) + new_key, new_value = _postprocess_single_tensor( + key, value, maxbound, quantization, is_modelopt_qlora + ) + if new_key is not None: + post_state_dict[new_key] = new_valueAlso applies to: 1050-1141
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/quant_utils.py` around lines 962 - 1049, Refactor postprocess_state_dict to delegate each tensor’s skip, replacement, KV-cache scaling, and scale-squeezing decisions to _postprocess_single_tensor instead of duplicating that control flow. Remove the separate RealQuant-scale and LoRA cleanup passes, pass the existing kv_cache_max_bound, kv_cache_format, and is_modelopt_qlora values through, and preserve tied-weight filtering and output assembly around the helper result.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 768-801: Handle the incompatibility between disk offload and
CPU-only execution in the setup flow that assigns device_map: either reject
_disk_offload when device is "cpu" with a clear error, or preserve
device_map="auto" so accelerate can honor max_memory and offload_folder. Ensure
the disk-offload branch no longer reports enabled while loading everything into
RAM.
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1581-1600: Add validation alongside the existing parser.error
checks in the argument-processing flow: when either max_cpu_memory_gb or
max_gpu_memory_gb is provided without offload_folder, call parser.error with a
clear message requiring --offload_folder. Preserve normal behavior when
offload_folder is set or neither budget flag is supplied.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1151-1186: Ensure every quantized module outside decoder layers is
processed before the resident parameter sweep: extend the export flow around
_dispatch_export_handler to discover remaining modules with an enabled weight
quantizer and dispatch each exactly once using a visited set, preserving
existing handling for decoder and offloaded modules. Do not allow such modules
to fall through to verbatim tensor streaming; if they cannot be dispatched, fail
explicitly instead.
- Around line 942-962: Update the fallback logic in _parse_shard_size so GB and
MB suffixes use decimal multipliers (1000**3 and 1000**2), matching
transformers.utils.convert_file_size_to_int. Keep GIB and MIB using binary
multipliers and preserve the existing parsing behavior for integers and
suffixless values.
- Around line 1054-1067: Update the streaming export flow in `_stream_tensor` to
apply tensor-level reverse conversion, including `Chunk` split rules, rather
than only using `build_reverse_name_mapper`; ensure converted tensors are
emitted under the correct unfused hub keys and shapes. If streaming cannot
support a mapping, fail fast instead of exporting an incorrect tensor, while
preserving the existing batch-path behavior.
In `@tests/unit/torch/export/test_offload_export.py`:
- Around line 274-286: Move the local imports of KV_CACHE_FP8, RealQuantLinear,
and sync_tied_input_amax to the file-level import section, then remove their
corresponding imports from the test functions. Keep the existing test behavior
and use the same internal module sources already referenced by those imports.
In `@tools/resource_monitor.py`:
- Line 58: Remove all three # nosec annotations from the subprocess import and
the nvidia-smi check_output and wrapped-command Popen usages in
resource_monitor.py. Preserve the existing fixed-argument, non-shell subprocess
behavior, and rely on codeowner review with PR justification rather than inline
Bandit suppression.
---
Nitpick comments:
In `@modelopt/torch/export/quant_utils.py`:
- Around line 962-1049: Refactor postprocess_state_dict to delegate each
tensor’s skip, replacement, KV-cache scaling, and scale-squeezing decisions to
_postprocess_single_tensor instead of duplicating that control flow. Remove the
separate RealQuant-scale and LoRA cleanup passes, pass the existing
kv_cache_max_bound, kv_cache_format, and is_modelopt_qlora values through, and
preserve tied-weight filtering and output assembly around the helper result.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1277-1283: Move the _has_accelerate_offload(model) validation to
the top of _export_transformers_checkpoint, immediately after dtype handling and
before PrepareMoEInputsRegistry dispatch or
requantize_resmooth_fused_llm_layers(model). Preserve the existing
NotImplementedError message and streaming-export guidance, and remove the later
duplicate check.
- Around line 877-889: The shard writer must clean up unfinalized temporary
parts when export fails. Add a close or context-manager cleanup path for the
writer containing _flush that unlinks every path in _part_files unless finalize
completed, and invoke it from a try/finally surrounding the streaming export
loops so failures before finalize cannot leave __shard_part_*.safetensors files
in the export directory.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d4fdbd00-3935-4f76-ac8e-0e77ea6d1d07
📒 Files selected for processing (19)
examples/hf_ptq/example_utils.pyexamples/hf_ptq/hf_ptq.pyexamples/hf_ptq/requirements.txtexamples/hf_ptq/scripts/huggingface_example.shmodelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/plugins/huggingface.pymodelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yamlmodelopt_recipes/ptq.mdtests/conftest.pytests/examples/hf_ptq/test_example_utils.pytests/gpu/torch/export/test_offload_export.pytests/tools/test_resource_monitor.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/quantization/plugins/test_fused_experts.pytools/resource_monitor.py
💤 Files with no reviewable changes (1)
- modelopt/torch/export/hf_export_handlers.py
| if _disk_offload: | ||
| for _k in max_memory: | ||
| if isinstance(_k, int): | ||
| if max_gpu_memory_gb is not None: | ||
| max_memory[_k] = int(max_gpu_memory_gb * 1024**3) | ||
| else: | ||
| max_memory[_k] = int(max_memory[_k] * gpu_mem_percentage) | ||
| if max_cpu_memory_gb is not None: | ||
| max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) | ||
| model_kwargs["max_memory"] = max_memory | ||
| print( | ||
| "Model does not fit to the GPU mem. " | ||
| f"We apply the following memory limit for calibration: \n{max_memory}\n" | ||
| "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " | ||
| "reduce the calibration `batch_size` manually." | ||
| "Disk-offload mode enabled. " | ||
| f"Memory budgets: {_fmt_max_memory(max_memory)}\n" | ||
| f"Offload folder: {offload_folder}\n" | ||
| "Weights exceeding GPU+CPU budgets will be streamed from disk." | ||
| ) | ||
| model_kwargs["max_memory"] = max_memory | ||
| else: | ||
| inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) | ||
| if "cpu" in inferred_device_map.values(): | ||
| for _device in max_memory: | ||
| if isinstance(_device, int): | ||
| max_memory[_device] *= gpu_mem_percentage | ||
|
|
||
| print( | ||
| "Model does not fit to the GPU mem. " | ||
| f"We apply the following memory limit for calibration: \n{max_memory}\n" | ||
| "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " | ||
| "reduce the calibration `batch_size` manually." | ||
| ) | ||
| model_kwargs["max_memory"] = max_memory | ||
|
|
||
| model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) | ||
| if _disk_offload: | ||
| model_kwargs2["offload_folder"] = offload_folder |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Disk-offload mode is a silent no-op when device="cpu".
device_map was already forced to "cpu" at Line 624-625. With device_map="cpu", accelerate ignores both max_memory and offload_folder, yet this branch still prints "Disk-offload mode enabled" and the GPU/CPU budgets. hf_ptq.py only rejects --low_memory_mode/--use_seq_device_map, so --offload_folder ... --device cpu reaches here and loads everything into RAM.
Either reject the combination up front or leave device_map="auto" when offloading.
🛠️ Suggested guard
if _disk_offload:
+ if device_map == "cpu":
+ raise ValueError(
+ "offload_folder requires device_map='auto'; it has no effect with device='cpu'."
+ )
for _k in max_memory:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if _disk_offload: | |
| for _k in max_memory: | |
| if isinstance(_k, int): | |
| if max_gpu_memory_gb is not None: | |
| max_memory[_k] = int(max_gpu_memory_gb * 1024**3) | |
| else: | |
| max_memory[_k] = int(max_memory[_k] * gpu_mem_percentage) | |
| if max_cpu_memory_gb is not None: | |
| max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) | |
| model_kwargs["max_memory"] = max_memory | |
| print( | |
| "Model does not fit to the GPU mem. " | |
| f"We apply the following memory limit for calibration: \n{max_memory}\n" | |
| "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " | |
| "reduce the calibration `batch_size` manually." | |
| "Disk-offload mode enabled. " | |
| f"Memory budgets: {_fmt_max_memory(max_memory)}\n" | |
| f"Offload folder: {offload_folder}\n" | |
| "Weights exceeding GPU+CPU budgets will be streamed from disk." | |
| ) | |
| model_kwargs["max_memory"] = max_memory | |
| else: | |
| inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) | |
| if "cpu" in inferred_device_map.values(): | |
| for _device in max_memory: | |
| if isinstance(_device, int): | |
| max_memory[_device] *= gpu_mem_percentage | |
| print( | |
| "Model does not fit to the GPU mem. " | |
| f"We apply the following memory limit for calibration: \n{max_memory}\n" | |
| "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " | |
| "reduce the calibration `batch_size` manually." | |
| ) | |
| model_kwargs["max_memory"] = max_memory | |
| model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) | |
| if _disk_offload: | |
| model_kwargs2["offload_folder"] = offload_folder | |
| if _disk_offload: | |
| if device_map == "cpu": | |
| raise ValueError( | |
| "offload_folder requires device_map='auto'; it has no effect with device='cpu'." | |
| ) | |
| for _k in max_memory: | |
| if isinstance(_k, int): | |
| if max_gpu_memory_gb is not None: | |
| max_memory[_k] = int(max_gpu_memory_gb * 1024**3) | |
| else: | |
| max_memory[_k] = int(max_memory[_k] * gpu_mem_percentage) | |
| if max_cpu_memory_gb is not None: | |
| max_memory["cpu"] = int(max_cpu_memory_gb * 1024**3) | |
| model_kwargs["max_memory"] = max_memory | |
| print( | |
| "Disk-offload mode enabled. " | |
| f"Memory budgets: {_fmt_max_memory(max_memory)}\n" | |
| f"Offload folder: {offload_folder}\n" | |
| "Weights exceeding GPU+CPU budgets will be streamed from disk." | |
| ) | |
| else: | |
| inferred_device_map = infer_auto_device_map(model, max_memory=max_memory) | |
| if "cpu" in inferred_device_map.values(): | |
| for _device in max_memory: | |
| if isinstance(_device, int): | |
| max_memory[_device] *= gpu_mem_percentage | |
| print( | |
| "Model does not fit to the GPU mem. " | |
| f"We apply the following memory limit for calibration: \n{max_memory}\n" | |
| "If you hit GPU OOM issue, please adjust `gpu_mem_percentage` or " | |
| "reduce the calibration `batch_size` manually." | |
| ) | |
| model_kwargs["max_memory"] = max_memory | |
| model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) | |
| if _disk_offload: | |
| model_kwargs2["offload_folder"] = offload_folder |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/example_utils.py` around lines 768 - 801, Handle the
incompatibility between disk offload and CPU-only execution in the setup flow
that assigns device_map: either reject _disk_offload when device is "cpu" with a
clear error, or preserve device_map="auto" so accelerate can honor max_memory
and offload_folder. Ensure the disk-offload branch no longer reports enabled
while loading everything into RAM.
| parser.add_argument( | ||
| "--max_cpu_memory_gb", | ||
| type=float, | ||
| default=None, | ||
| help=( | ||
| "Maximum CPU RAM budget in GiB for disk-offload model loading. " | ||
| "Only effective when --offload_folder is set. " | ||
| "Weights beyond this limit are streamed from disk." | ||
| ), | ||
| ) | ||
| parser.add_argument( | ||
| "--max_gpu_memory_gb", | ||
| type=float, | ||
| default=None, | ||
| help=( | ||
| "Maximum GPU memory budget per device in GiB for disk-offload model loading. " | ||
| "Only effective when --offload_folder is set. " | ||
| "Defaults to 80%% of available GPU memory when not specified." | ||
| ), | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Memory budgets are silently ignored without --offload_folder.
Both flags are no-ops unless --offload_folder is set (as the help text admits), but nothing tells the user. Add a check next to the other parser.error validations.
🛠️ Suggested validation
+ if args.offload_folder is None and (
+ args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None
+ ):
+ parser.error(
+ "--max_cpu_memory_gb/--max_gpu_memory_gb require --offload_folder (disk-offload mode)."
+ )
+
return args📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| parser.add_argument( | |
| "--max_cpu_memory_gb", | |
| type=float, | |
| default=None, | |
| help=( | |
| "Maximum CPU RAM budget in GiB for disk-offload model loading. " | |
| "Only effective when --offload_folder is set. " | |
| "Weights beyond this limit are streamed from disk." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--max_gpu_memory_gb", | |
| type=float, | |
| default=None, | |
| help=( | |
| "Maximum GPU memory budget per device in GiB for disk-offload model loading. " | |
| "Only effective when --offload_folder is set. " | |
| "Defaults to 80%% of available GPU memory when not specified." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--max_gpu_memory_gb", | |
| type=float, | |
| default=None, | |
| help=( | |
| "Maximum GPU memory budget per device in GiB for disk-offload model loading. " | |
| "Only effective when --offload_folder is set. " | |
| "Defaults to 80%% of available GPU memory when not specified." | |
| ), | |
| ) | |
| if args.offload_folder is None and ( | |
| args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None | |
| ): | |
| parser.error( | |
| "--max_cpu_memory_gb/--max_gpu_memory_gb require --offload_folder (disk-offload mode)." | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/hf_ptq.py` around lines 1581 - 1600, Add validation alongside
the existing parser.error checks in the argument-processing flow: when either
max_cpu_memory_gb or max_gpu_memory_gb is provided without offload_folder, call
parser.error with a clear message requiring --offload_folder. Preserve normal
behavior when offload_folder is set or neither budget flag is supplied.
| def _parse_shard_size(size: int | str) -> int: | ||
| """Convert a shard-size string (e.g. ``"10GB"``, ``"500MB"``) to bytes.""" | ||
| try: | ||
| from transformers.utils import convert_file_size_to_int | ||
|
|
||
| return convert_file_size_to_int(size) | ||
| except ImportError: | ||
| pass | ||
| if isinstance(size, int): | ||
| return size | ||
| s = size.strip().upper() | ||
| if s.endswith("GIB"): | ||
| return int(float(s[:-3]) * 1024**3) | ||
| if s.endswith("GB"): | ||
| return int(float(s[:-2]) * 1024**3) | ||
| if s.endswith("MIB"): | ||
| return int(float(s[:-3]) * 1024**2) | ||
| if s.endswith("MB"): | ||
| return int(float(s[:-2]) * 1024**2) | ||
| return int(s) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
transformers convert_file_size_to_int GB GiB bytes conversion
💡 Result:
The function convert_file_size_to_int, commonly found in Hugging Face libraries like transformers, datasets, and accelerate, converts a string representing a file size into an integer in bytes [1][2][3]. The function supports both binary units (e.g., KiB, MiB, GiB) and decimal units (e.g., KB, MB, GB), as well as their bit-based counterparts [1][2][4]. The logic used for these conversions is as follows: Binary Units (Powers of 1024): - KiB: 2^10 (1,024) [1][2][4] - MiB: 2^20 (1,048,576) [1][2][4] - GiB: 2^30 (1,073,741,824) [1][2][4] - TiB/PiB: Often supported in some implementations (e.g., datasets) as 2^40 and 2^50 respectively [2][4]. Decimal Units (Powers of 10): - KB: 10^3 (1,000) [1][2][4] - MB: 10^6 (1,000,000) [1][2][4] - GB: 10^9 (1,000,000,000) [1][2][4] Bit-based variants: - For decimal units (KB, MB, GB), the function detects if the string ends in 'b' (e.g., "5Kb"). If it does, the calculated value is divided by 8 to convert from bits to bytes [1][2][4]. Note that implementations may vary slightly between libraries (e.g., some versions use integer multiplication while others cast to float before multiplying) [1][2][3]. The function is designed to raise a ValueError if the input string format is unrecognized [1][2][3].
Citations:
- 1: https://github.com/huggingface/transformers/blob/c472755e/src/transformers/utils/hub.py
- 2: https://github.com/huggingface/datasets/blob/master/src/datasets/utils/py_utils.py
- 3: https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/modeling.py
- 4: https://github.com/huggingface/datasets/blob/main/src/datasets/utils/py_utils.py
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the local implementation around the helper and its callers.
FILE="modelopt/torch/export/unified_export_hf.py"
wc -l "$FILE"
sed -n '920,980p' "$FILE"
# Search for other shard-size parsing helpers or related usages.
rg -n "_parse_shard_size|convert_file_size_to_int|shard size|shard_size" modelopt/torch/export/unified_export_hf.py modelopt -SRepository: NVIDIA/Model-Optimizer
Length of output: 7328
Align the fallback shard-size units with transformers
When transformers isn’t installed, GB/MB are treated as binary here (1024**3 / 1024**2), but convert_file_size_to_int uses decimal units for those suffixes. That makes the fallback allow slightly larger shards than intended; mirror the upstream decimal semantics.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/export/unified_export_hf.py` around lines 942 - 962, Update
the fallback logic in _parse_shard_size so GB and MB suffixes use decimal
multipliers (1000**3 and 1000**2), matching
transformers.utils.convert_file_size_to_int. Keep GIB and MIB using binary
multipliers and preserve the existing parsing behavior for integers and
suffixless values.
| def test_postprocess_kv_scale_renamed_and_divided(): | ||
| """k_bmm_quantizer._amax is renamed to k_proj.k_scale and divided by maxbound.""" | ||
| from modelopt.torch.export.model_config import KV_CACHE_FP8 | ||
|
|
||
| key, val = _postprocess_single_tensor( | ||
| "model.layers.0.self_attn.k_bmm_quantizer._amax", | ||
| torch.tensor(224.0), | ||
| 448.0, | ||
| KV_CACHE_FP8, | ||
| ) | ||
| assert key == "model.layers.0.self_attn.k_proj.k_scale" | ||
| assert abs(val.item() - 0.5) < 1e-5 | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Move in-function imports to the top of the file.
KV_CACHE_FP8 (L276), RealQuantLinear (L298), and sync_tied_input_amax (L351) are each imported locally inside their test function. None are circular-import or optional-dependency cases — all are plain internal modules, and _postprocess_single_tensor from the same quant_utils module is already imported at the top of the file. As per path instructions, "Imports inside functions or test methods without explicit justification" should be flagged, since "Imports belong at the top of the file so import errors surface at collection time, not mid-test."
🔧 Proposed fix: hoist imports to the top
import modelopt.torch.quantization as mtq
+from modelopt.torch.export.model_config import KV_CACHE_FP8
from modelopt.torch.export.quant_utils import _postprocess_single_tensor
+from modelopt.torch.export.quant_utils import sync_tied_input_amax
from modelopt.torch.export.unified_export_hf import (
_export_quantized_weight,
_has_accelerate_offload,
_StreamingShardWriter,
)
+from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinearThen remove the corresponding local from ... import ... lines at L276, L298, and L351.
Also applies to: 296-303, 345-364
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/torch/export/test_offload_export.py` around lines 274 - 286, Move
the local imports of KV_CACHE_FP8, RealQuantLinear, and sync_tied_input_amax to
the file-level import section, then remove their corresponding imports from the
test functions. Keep the existing test behavior and use the same internal module
sources already referenced by those imports.
Source: Path instructions
| import csv | ||
| import shutil | ||
| import signal | ||
| import subprocess # nosec B404 - wraps and monitors a user-provided workload command; no shell is used |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Remove # nosec bypasses; get codeowner sign-off instead.
Three # nosec annotations suppress Bandit here (import, nvidia-smi check_output, and the wrapped-command Popen). Per repo policy, inline # nosec is not an allowed bypass regardless of the justification comment attached — a security-sensitive pattern that's genuinely necessary must go through @NVIDIA/modelopt-setup-codeowners review with justification in the PR description, not an inline suppression.
The underlying subprocess usage itself looks safe (fixed argv lists, no shell=True), but the suppression mechanism is what's disallowed.
As per coding guidelines: "Bandit security checks must pass without exceptions. # nosec comments are not allowed as a bypass for security checks." Also: "Any use of '# nosec' comments to bypass Bandit security checks is not allowed. If a security-sensitive pattern is genuinely necessary, the PR must be reviewed and approved by @NVIDIA/modelopt-setup-codeowners with an explicit justification in the PR description."
Also applies to: 165-172, 361-361
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tools/resource_monitor.py` at line 58, Remove all three # nosec annotations
from the subprocess import and the nvidia-smi check_output and wrapped-command
Popen usages in resource_monitor.py. Preserve the existing fixed-argument,
non-shell subprocess behavior, and rely on codeowner review with PR
justification rather than inline Bandit suppression.
Source: Coding guidelines
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 3
♻️ Duplicate comments (3)
examples/hf_ptq/example_utils.py (1)
768-801: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDisk-offload mode is still a silent no-op with
--device cpu.
device_mapis forced to"cpu"at Line 624-625, and accelerate ignoresmax_memory/offload_folderfor a plain"cpu"device map — yet this branch prints "Disk-offload mode enabled" and everything loads into RAM. Reject the combination (here or inparse_args) or keepdevice_map="auto"whenoffload_folderis set.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/example_utils.py` around lines 768 - 801, Prevent disk-offload mode from being used with a forced device_map of "cpu": either reject this argument combination during argument validation or preserve device_map="auto" whenever offload_folder is configured. Update the disk-offload branch around _disk_offload so it cannot claim disk offloading while loading entirely into RAM.modelopt/torch/export/unified_export_hf.py (2)
942-961: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winFallback shard-size units still differ from
transformers.
GB/MBare treated as binary here whileconvert_file_size_to_inttreats them as decimal, so the no-transformers fallback permits larger shards than requested. Use1000**3/1000**2for the non-iBsuffixes.transformers convert_file_size_to_int GB MB decimal or binary units🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 942 - 961, Update _parse_shard_size so non-binary GB and MB suffixes use decimal conversion with 1000**3 and 1000**2 respectively, matching transformers.convert_file_size_to_int; keep GIB and MIB conversions binary with 1024-based factors.
1151-1186: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftQuantized modules that are resident and unhooked still bypass handler dispatch.
Dispatch runs only for decoder-layer submodules (Line 1105-1111) and non-decoder modules that both carry an offload hook and currently hold meta tensors (Line 1155-1163). A quantized module outside the decoder stack that happens to be resident (GPU-placed
lm_head, MTP/draft submodule, vision-tower linear on mixed placement) falls through to the resident sweep at Line 1179-1186 and gets its unpackedweightstreamed whileget_quant_configadvertises it as quantized. Dispatch remaining modules with an enabled weight quantizer (guarded by a visited set), or raise when one reaches the resident sweep.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 1151 - 1186, Update the export traversal around _dispatch_export_handler and the resident parameter sweep so resident, unhooked quantized modules are not streamed as unpacked weights. Track visited modules and dispatch every module with an enabled weight quantizer, including non-decoder modules, while avoiding duplicate dispatch; alternatively, make the resident sweep raise when it encounters an undispatched quantized module.
🧹 Nitpick comments (3)
modelopt/torch/export/unified_export_hf.py (2)
1139-1149: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument that the streaming export leaves the model unusable.
Nulling
_buffers/_parametersentries is necessary for memory, but the model is silently left structurally broken afterwards. Worth stating in the_export_transformers_checkpoint_streamingdocstring so callers don't reuse the model post-export.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 1139 - 1149, Update the _export_transformers_checkpoint_streaming docstring to explicitly state that streaming export nulls CUDA buffers and parameters and leaves the model structurally unusable afterward, so callers must not reuse the model.
891-910: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAlias check is a linear scan per
add().
any(ptr == tensor.data_ptr() for ptr, _, _ in self._buffer_storage)is O(buffered tensors) on every call, so filling a 10 GB shard with many small tensors is quadratic. Keep aset[int]of buffered pointers alongside_buffer_storage(reset in_flush) for O(1) lookup.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 891 - 910, Update the tensor buffering logic in add to maintain a dedicated set of buffered data pointers alongside _buffer_storage, using constant-time membership checks instead of scanning _buffer_storage. Add each original tensor pointer when buffering, and clear the pointer set in _flush together with the existing buffer state.examples/hf_ptq/example_utils.py (1)
726-730: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueComment contradicts the branch it documents.
The comment says the built-in class is what the offload/streaming paths were validated against, but
use_bundled_codemakestrust_remote_code=Trueselect the bundled code path. Reword to state the intent (remote-code trust opts into bundled DeepSeek modeling code; otherwise use the built-in class).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/hf_ptq/example_utils.py` around lines 726 - 730, Update the comment above use_bundled_code and the following architecture branch to accurately state that trust_remote_code enables bundled DeepSeek modeling code, while the default path uses the built-in class. Keep the implementation unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1622-1631: Extend the argument validation near the existing
offload_folder checks to reject max_cpu_memory_gb or max_gpu_memory_gb when
offload_folder is unset, since those budgets only apply to disk offload. Also
reject combining offload_folder with --device cpu, following the constraint
documented by example_utils.get_model.
In `@modelopt/torch/export/quant_utils.py`:
- Line 1039: Update the KV-cache warning condition in the export path around
kv_cache_format and value so it does not call value.item() while the tensor is
still device-resident. Move the scalar comparison until after value has been
transferred to CPU, or batch it with the existing export-side warning checks,
while preserving the warning threshold and behavior.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1092-1102: The streaming path in _stream_tensor must handle
tensor-level split rules that batch export processes via
revert_weight_conversion_quant_aware(), rather than applying only name_mapper
renames. Apply the equivalent reverse weight conversion before writer.add, or
fail fast when the mapping contains any non-rename rules such as Chunk; preserve
existing alias filtering and tensor serialization.
---
Duplicate comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 768-801: Prevent disk-offload mode from being used with a forced
device_map of "cpu": either reject this argument combination during argument
validation or preserve device_map="auto" whenever offload_folder is configured.
Update the disk-offload branch around _disk_offload so it cannot claim disk
offloading while loading entirely into RAM.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 942-961: Update _parse_shard_size so non-binary GB and MB suffixes
use decimal conversion with 1000**3 and 1000**2 respectively, matching
transformers.convert_file_size_to_int; keep GIB and MIB conversions binary with
1024-based factors.
- Around line 1151-1186: Update the export traversal around
_dispatch_export_handler and the resident parameter sweep so resident, unhooked
quantized modules are not streamed as unpacked weights. Track visited modules
and dispatch every module with an enabled weight quantizer, including
non-decoder modules, while avoiding duplicate dispatch; alternatively, make the
resident sweep raise when it encounters an undispatched quantized module.
---
Nitpick comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 726-730: Update the comment above use_bundled_code and the
following architecture branch to accurately state that trust_remote_code enables
bundled DeepSeek modeling code, while the default path uses the built-in class.
Keep the implementation unchanged.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1139-1149: Update the _export_transformers_checkpoint_streaming
docstring to explicitly state that streaming export nulls CUDA buffers and
parameters and leaves the model structurally unusable afterward, so callers must
not reuse the model.
- Around line 891-910: Update the tensor buffering logic in add to maintain a
dedicated set of buffered data pointers alongside _buffer_storage, using
constant-time membership checks instead of scanning _buffer_storage. Add each
original tensor pointer when buffering, and clear the pointer set in _flush
together with the existing buffer state.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b615ba93-877e-467a-b68b-5db797ebc9d5
📒 Files selected for processing (14)
examples/hf_ptq/example_utils.pyexamples/hf_ptq/hf_ptq.pymodelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/plugins/huggingface.pymodelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yamlmodelopt_recipes/ptq.mdtests/examples/hf_ptq/test_example_utils.pytests/gpu/torch/export/test_offload_export.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/quantization/plugins/test_fused_experts.py
💤 Files with no reviewable changes (1)
- modelopt/torch/export/hf_export_handlers.py
🚧 Files skipped from review as they are similar to previous changes (2)
- modelopt/torch/quantization/plugins/huggingface.py
- modelopt_recipes/ptq.md
| if args.offload_folder is not None and args.low_memory_mode: | ||
| parser.error("--offload_folder (disk-offload) is not compatible with --low_memory_mode.") | ||
|
|
||
| if args.offload_folder is not None and args.use_seq_device_map: | ||
| parser.error( | ||
| "--offload_folder (disk-offload) is not compatible with --use_seq_device_map; " | ||
| "device_map=auto is used for disk-offload to let accelerate place layers across " | ||
| "GPU, CPU, and disk." | ||
| ) | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Budget flags are still accepted without --offload_folder.
--max_cpu_memory_gb / --max_gpu_memory_gb are no-ops unless --offload_folder is set (per their own help text) but nothing rejects that combination. Same place is the natural spot to also reject --offload_folder with --device cpu (see the example_utils.get_model comment).
🛠️ Suggested validation
+ if args.offload_folder is None and (
+ args.max_cpu_memory_gb is not None or args.max_gpu_memory_gb is not None
+ ):
+ parser.error(
+ "--max_cpu_memory_gb/--max_gpu_memory_gb require --offload_folder (disk-offload mode)."
+ )
+
return args🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/hf_ptq.py` around lines 1622 - 1631, Extend the argument
validation near the existing offload_folder checks to reject max_cpu_memory_gb
or max_gpu_memory_gb when offload_folder is unset, since those budgets only
apply to disk offload. Also reject combining offload_folder with --device cpu,
following the constraint documented by example_utils.get_model.
| ) | ||
| assert kv_cache_max_bound > 0, "Maxbound must be greater than zero." | ||
| value = value.float() / kv_cache_max_bound | ||
| if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5: |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline modelopt/torch/export/quant_utils.py --view expanded | sed -n '1,240p'
printf '\n--- slice around reported lines ---\n'
sed -n '1000,1060p' modelopt/torch/export/quant_utils.py
printf '\n--- nearby references to kv_cache_format and value ---\n'
rg -n "kv_cache_format|value\.item\(|value\b" modelopt/torch/export/quant_utils.pyRepository: NVIDIA/Model-Optimizer
Length of output: 7807
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("modelopt/torch/export/quant_utils.py")
lines = p.read_text().splitlines()
for start, end in [(1015, 1055), (920, 980)]:
print(f"\n--- {p}:{start}-{end} ---")
for i in range(start-1, min(end, len(lines))):
print(f"{i+1:5d}: {lines[i]}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 4885
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- uses of _postprocess_single_tensor ---'
rg -n "_postprocess_single_tensor\(" -S modelopt/torch/export/quant_utils.py modelopt || true
printf '\n%s\n' '--- streaming/state-dict path around postprocess_state_dict ---'
sed -n '1050,1145p' modelopt/torch/export/quant_utils.py
printf '\n%s\n' '--- any explicit cpu()/to(\"cpu\") around this path ---'
rg -n "to\\(['\"]cpu['\"]\\)|\\.cpu\\(" modelopt/torch/export/quant_utils.pyRepository: NVIDIA/Model-Optimizer
Length of output: 4390
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1060,1115p' modelopt/torch/export/unified_export_hf.py
printf '\n%s\n' '--- surrounding helpers/entrypoints ---'
rg -n "stream|state_dict|_postprocess_single_tensor|postprocess_state_dict" modelopt/torch/export/unified_export_hf.pyRepository: NVIDIA/Model-Optimizer
Length of output: 7192
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- enable_weight_access_and_writeback definition ---'
rg -n "def enable_weight_access_and_writeback|class enable_weight_access_and_writeback|enable_weight_access_and_writeback\(" modelopt -S
printf '\n%s\n' '--- streaming export around the offload context ---'
sed -n '1088,1188p' modelopt/torch/export/unified_export_hf.py
printf '\n%s\n' '--- surrounding offload helper docs ---'
sed -n '1,260p' modelopt/torch/export/unified_export_hf.pyRepository: NVIDIA/Model-Optimizer
Length of output: 19523
Avoid the .item() sync in the streaming KV-cache warning path. This branch runs before the tensor is moved to CPU, so value.item() adds a device-host sync just to decide whether to log. Move the check after .cpu() or batch it with other export-side warnings.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/export/quant_utils.py` at line 1039, Update the KV-cache
warning condition in the export path around kv_cache_format and value so it does
not call value.item() while the tensor is still device-resident. Move the scalar
comparison until after value has been transferred to CPU, or batch it with the
existing export-side warning checks, while preserving the warning threshold and
behavior.
Source: Coding guidelines
| def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: | ||
| new_key, new_value = _postprocess_single_tensor( | ||
| full_key, tensor, kv_cache_max_bound, kv_cache_format, is_modelopt_qlora | ||
| ) | ||
| if new_key is None or new_value is None: | ||
| return | ||
| if name_mapper is not None: | ||
| new_key = name_mapper(new_key) | ||
| if new_key in tied_alias_keys: | ||
| return | ||
| writer.add(new_key, new_value.detach().contiguous().cpu()) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Streaming path still applies rename-only reverse conversion.
build_reverse_name_mapper() covers renames only; the batch path additionally runs revert_weight_conversion_quant_aware() for tensor-level split (Chunk) rules. A model whose conversion mapping includes a split will emit a fused tensor under an unfused hub key. Either apply the split in _stream_tensor or fail fast when the mapping contains non-rename rules.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/export/unified_export_hf.py` around lines 1092 - 1102, The
streaming path in _stream_tensor must handle tensor-level split rules that batch
export processes via revert_weight_conversion_quant_aware(), rather than
applying only name_mapper renames. Apply the equivalent reverse weight
conversion before writer.add, or fail fast when the mapping contains any
non-rename rules such as Chunk; preserve existing alias filtering and tensor
serialization.
The streaming exporter decided which modules to materialize by testing accelerate internals directly (_hf_hook, _get_offload_hook, meta params), so it only ever worked for accelerate offload even though the mechanism it drives -- enable_weight_access_and_writeback -- already dispatches over FSDP2, HF TP and accelerate alike. Add requires_weight_materialization() beside that dispatcher, mirroring its branches, and have the exporter ask it instead. Two conditions must hold: the module owns tensors that are not readable right now (meta, or a sharded DTensor), and a context exists that can materialize them. Also pass the real root model and a cached name_to_module to both windows -- FSDP2 detection walks up from root_model, so passing the layer as its own root could never detect sharding, and the cache avoids the O(N^2) rescan the parameter was added for. Fixes a bug this surfaced: the non-decoder pass skipped decoder layers by id but not their children. An offloaded layer's children return to meta when its window closes, so they passed the filter and were exported a second time, hitting a None amax on weights the first pass had already packed. Only reproduced when the device map placed a layer on CPU, which is why the GPU tests caught it and the unit tests did not. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
b2d731d to
3175d37
Compare
|
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (2)
modelopt/torch/quantization/utils/core_utils.py (1)
623-647: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a brief reason for the deferred import.
from ..plugins.accelerate import _get_offload_hookis function-local. As per coding guidelines, "use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment" — a one-line comment (optionalacceleratedependency) would make the intent explicit here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/utils/core_utils.py` around lines 623 - 647, Add a brief inline comment immediately before the function-local import of _get_offload_hook in requires_weight_materialization, explaining that the local import avoids eagerly loading the optional Accelerate dependency. Keep the import and surrounding logic unchanged.Source: Coding guidelines
modelopt/torch/export/unified_export_hf.py (1)
1281-1295: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRaise the offload
NotImplementedErrorbefore the model is mutated.
requantize_resmooth_fused_llm_layers(model)(Line 1279) and the MoE input preparation above it already mutated the model by the time this guard fires, so a caller that catches the error is left with a partially processed model. Hoisting the check to the top of the function makes the failure non-destructive.♻️ Suggested reordering
def _export_transformers_checkpoint( model: nn.Module, dtype: torch.dtype | None = None, is_modelopt_qlora: bool = False, **kwargs, ) -> tuple[dict[str, Any], dict[str, Any]]: + # Offloaded models need their weights materialized layer-by-layer, which this + # whole-state-dict path cannot do; export_hf_checkpoint() streams them instead. + # Checked before any model mutation so the failure leaves the model untouched. + if _has_accelerate_offload(model): + raise NotImplementedError( + "_export_transformers_checkpoint does not support disk/CPU-offloaded models. " + "Use export_hf_checkpoint() which dispatches to " + "_export_transformers_checkpoint_streaming." + )and drop the block at Lines 1281-1287.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/unified_export_hf.py` around lines 1281 - 1295, Move the _has_accelerate_offload(model) guard to the beginning of _export_transformers_checkpoint, before requantize_resmooth_fused_llm_layers(model) and MoE input preparation mutate the model. Preserve the existing NotImplementedError message, and remove the later duplicate guard while leaving hook removal unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/hf_ptq/example_utils.py`:
- Around line 686-694: The model-loading branches for is_speculative,
has_pack_quantized_config, and mxfp4 bypass _disk_offload, so offload_folder and
max_cpu_memory_gb are ignored. Update those from_pretrained paths to honor disk
offload consistently, or explicitly warn/error at the entry point when the
selected checkpoint type cannot support it.
---
Nitpick comments:
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1281-1295: Move the _has_accelerate_offload(model) guard to the
beginning of _export_transformers_checkpoint, before
requantize_resmooth_fused_llm_layers(model) and MoE input preparation mutate the
model. Preserve the existing NotImplementedError message, and remove the later
duplicate guard while leaving hook removal unchanged.
In `@modelopt/torch/quantization/utils/core_utils.py`:
- Around line 623-647: Add a brief inline comment immediately before the
function-local import of _get_offload_hook in requires_weight_materialization,
explaining that the local import avoids eagerly loading the optional Accelerate
dependency. Keep the import and surrounding logic unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6bc48131-9f59-4270-88f1-7d64a43ffbfa
📒 Files selected for processing (15)
examples/hf_ptq/example_utils.pyexamples/hf_ptq/hf_ptq.pymodelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/quantization/plugins/huggingface.pymodelopt/torch/quantization/utils/core_utils.pymodelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yamlmodelopt_recipes/ptq.mdtests/examples/hf_ptq/test_example_utils.pytests/gpu/torch/export/test_offload_export.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/quantization/plugins/test_fused_experts.py
💤 Files with no reviewable changes (1)
- modelopt/torch/export/hf_export_handlers.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt_recipes/ptq.md
| _disk_offload = offload_folder is not None | ||
| if _disk_offload and max_cpu_memory_gb is None: | ||
| warnings.warn( | ||
| "offload_folder is set but max_cpu_memory_gb is not specified. " | ||
| "CPU memory usage during model load will be unbounded. " | ||
| "Pass max_cpu_memory_gb to cap CPU usage.", | ||
| UserWarning, | ||
| ) | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
offload_folder is silently ignored for speculative / pack-quantized / MXFP4 checkpoints.
_disk_offload is only consumed inside the final else branch (Lines 840-873). The is_speculative, has_pack_quantized_config, and mxfp4 branches above build their own from_pretrained calls and never receive max_memory or offload_folder, so those checkpoints load fully resident while the caller believes disk-offload is active. Since this warning already fires at the entry point, consider also warning (or erroring) here when disk-offload is requested for a load path that cannot honor it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/hf_ptq/example_utils.py` around lines 686 - 694, The model-loading
branches for is_speculative, has_pack_quantized_config, and mxfp4 bypass
_disk_offload, so offload_folder and max_cpu_memory_gb are ignored. Update those
from_pretrained paths to honor disk offload consistently, or explicitly
warn/error at the entry point when the selected checkpoint type cannot support
it.
5 participants
What does this PR do?
Type of change: New feature, bug fix, new tests
Enables single-GPU PTQ for models too large to fit in VRAM (e.g. Nemotron-Ultra-550B at 1.1 TB BF16, DeepSeek-R1 at 642 GB BF16) by adding accelerate disk/CPU offload support to the HF PTQ example and fixing the export path to correctly handle offloaded models.
G1 Offload-aware unified HF export (
modelopt/torch/export/unified_export_hf.py)The existing
_export_transformers_checkpointremoved accelerate hooks before materializing weights, silently writing meta tensors (empty weights) to the checkpoint. Fix:_has_accelerate_offload(model)detects any disk/CPU-offload accelerate hook in the model tree._process_quantized_modules_offloaded(model, dtype)new export path for offloaded models: materializes one decoder layer at a time viaenable_weight_access_and_writeback, dispatches export handlers inside the context window, and snapshots the layer state dict before hooks re-offload the weights. A second pass collects non-decoder modules that are also disk-offloaded (embed, norm, lm_head) to avoid meta tensors in the returned state dict. Hooks are removed only after the full state dict is assembled._export_quantized_weightraisesRuntimeErroron meta input instead of silently corrupting the checkpoint.G2 Disk-offload CLI (
examples/hf_ptq/hf_ptq.py,example_utils.py)Three new arguments to
hf_ptq.py:--offload_folder PATHenable accelerate disk offload; shards spill here.--max_gpu_memory_gb NVRAM budget for the accelerate device map.--max_cpu_memory_gb NCPU RAM budget for the accelerate device map.Validation:
--offload_folderis incompatible with--low_memory_modeand--use_seq_device_map.G3 Streaming shard writer for 80 GB CPU RAM (
modelopt/torch/export/unified_export_hf.py)The G1 path accumulated the entire quantized state dict in CPU RAM before writing
(~764 GiB for Ultra 550B), blocking the 80 GB target.
New streaming path writes shard files layer-by-layer. Peak memory = 1 decoder layer +
1 shard buffer instead of the full checkpoint:
Key pieces:
_StreamingShardWriter(export_dir, max_shard_size)buffers tensors up tomax_shard_sizebytes, flushes to numbered temp files (__shard_part_NNNNN.safetensors), renames to canonical shard names atfinalize(), writesmodel.safetensors.index.json. Single-shard exports producemodel.safetensorswith no index file._postprocess_single_tensor(key, value, ...)per-tensor extraction ofpostprocess_state_dictlogic (KV amax scale, skip/rename, squeeze) for streaming use._parse_shard_size(size)converts"10GB"/"500MB"strings to bytes._export_transformers_checkpoint_streaming(model, dtype, export_dir, max_shard_size)streams decoder layers viaenable_weight_access_and_writeback, applies per-tensor postprocessing + name reversal + tied-alias filter, writes shard files directly. Non-decoder offloaded modules and GPU-resident tensors are handled in separate passes.export_hf_checkpointdispatches to the streaming path when_has_accelerate_offload(model)is true;hf_quant_config.json, quant-config name reversal, andconfig.jsonupdate are shared between paths.export_hf_checkpointaccepts a newmax_shard_sizeparameter (default"10GB") that controls the shard size for both paths.Supporting changes
modelopt/torch/quantization/plugins/huggingface.pyget_nemotron_h_decoder_layersnow checks bothmodel.backbone.layers(remote-code variant) andmodel.model.layers(native HF variant), fixing layer discovery for NemotronH when loaded withouttrust_remote_code.modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yamlnew recipe combining NVFP4 W4A4 on MoE experts, FP8 KV cache, and layerwise calibration withcalib_mutates_weights: false(required for disk-offload compatibility).example_utils.py_FP8BF16Fallbackshim: dequantizes block-scaled FP8 expert weights to BF16 for calibration forward passes when thekernelspackage is unavailable (e.g. DSR1 on nodes without finegrained FP8 kernel support).Usage
Testing
Unit tests (
tests/unit/torch/export/test_offload_export.py, 7 tests, CPU-only):_has_accelerate_offloaddetection (true/false/nested-module cases)_export_quantized_weightmeta-tensor guard (raises on meta, passes on real)_process_quantized_modules_offloadedwith disk-offloaded embed + GPU-resident decoder layer: verifies no meta tensor in returned state dictGPU integration tests (
tests/gpu/torch/export/test_offload_export.py, 2 tests):hf_quant_config.jsoncalib_mutates_weights=False): disk-offload path end-to-endValidation runs (on GB200, single GPU):
AttributeError: NemotronHForCausalLM has no attribute 'backbone'â fixed by the non-decoder offloaded tensor materialization pass.kernelspackage dependency.End-to-end validation
Two production-scale checkpoints were quantized end-to-end using the new disk-offload PTQ path on a single GB200 GPU (189 GiB VRAM).
DeepSeek-R1 (671B, MoE)
DeepseekV3ForCausalLM, 671B params, 61 decoder layersnvfp4_experts_only-kv_fp8_layerwise_offload--max_gpu_memory_gb--max_cpu_memory_gb--calib_size/--batch_size--trust_remote_codeNVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 (550B, NemotronH MoE + Mamba)
NemotronHForCausalLM, ~550B params, 108 decoder layersnvfp4_experts_only-kv_fp8_layerwise_offload--max_gpu_memory_gb/--max_cpu_memory_gb--calib_size/--batch_size--trust_remote_codeNemotronHForCausalLM)Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅Additional Information
Summary by CodeRabbit